
import sys
import tensorflow as tf
import tensorflow.keras.utils
import pandas as pd


def get_compiled_model():
    model = tf.keras.Sequential([
        tf.keras.layers.Dense(10, activation='relu'),
        tf.keras.layers.Dense(10, activation='relu'),
        tf.keras.layers.Dense(1, activation='sigmoid')
    ])

    model.compile(optimizer='adam',
                  loss='binary_crossentropy',
                  metrics=['accuracy'])
    return model


def data_parse(file):
    data = pd.read_csv(file)
    data = data.sort_values(by='trade_date', ascending=True)
    data = pd.DataFrame(data, columns=['open','close','high','low'])
    # x = [int(str(x)[0:6]) for x in data['trade_date']]
    # data['ts_code'] = x
    print(data.dtypes)
    print(data.head())
    jine = data.copy()
    jine['y'] = jine.apply(lambda x: sum(x)/4 > 25, axis=1)
    print(jine.head())
    target = jine.pop('y')
    print(target.dtypes)
    td = tf.data.Dataset.from_tensor_slices((data.values, target.values))
    for k, v in td.take(10):
        print("{}, {}".format(k, v))

    train_td = td.take(2000)
    valid_td = td.skip(2000)
    train_data = train_td.batch(100)
    valid_data = valid_td.batch(100)
    train, train_label = next(iter(train_data))
    print(train[0], train_label[0])

def train_it(data, vdata):
    model = get_compiled_model()
    model.fit(data, epochs=5)

    model.evaluate(vdata)

if __name__ == '__main__':
    data_parse(sys.argv[1])
